import ctypes
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt
import time


# 定义结构体,这里结构体顺序必须和打包dll时相同
class Detection(ctypes.Structure):
    _fields_ = [
        ('conf', ctypes.c_float),
        ('clas_id', ctypes.c_size_t),
        # x1,y1,w,h
        ('box_array', ctypes.POINTER(ctypes.c_int)),
    ]


class MyClass(object):
    def __init__(self, engine_file, dll_path, class_name, conf_threshold):
        self.my_class_dll = ctypes.cdll.LoadLibrary(dll_path)
        # self.my_class_dll = ctypes.CDLL(r"./lib/mylib.dll")
        # self.my_class_dll = ctypes.PyDLL(r"./lib/mylib.dll")
        create_my_class = self.my_class_dll.InitDetect
        # 输入值类型
        create_my_class.argtypes = [ctypes.c_char_p, ctypes.POINTER(ctypes.c_float)]
        # 返回值类型,构造函数返回一个指针
        create_my_class.restype = ctypes.POINTER(ctypes.c_void_p)

        conf_threshold = ctypes.c_float(conf_threshold)

        self.obj = create_my_class(engine_file.encode(), conf_threshold)
        self.class_name = class_name
        num_classes = len(class_name)
        self.color_list = self.SetColor(num_classes)

    def OutResult(self, img):
        my_class_out_result = self.my_class_dll.OutResult
        my_class_out_result.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.POINTER(ctypes.c_int),
                                        ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_ubyte),
                                        ctypes.POINTER(ctypes.c_size_t)]
        my_class_out_result.restype = ctypes.POINTER(Detection)

        self.result_size = ctypes.c_size_t()
        height = ctypes.c_int(img.shape[0])
        width = ctypes.c_int(img.shape[1])
        channel = ctypes.c_int(img.shape[2])
        # 将图像数据转为uchar指针数组形式
        img_data = img.ctypes.data_as(ctypes.POINTER(ctypes.c_ubyte))

        self.result_array = my_class_out_result(self.obj, height, width, img_data, self.result_size)
        result_size = self.result_size.value
        result_box = self.Array2Box(self.result_array, result_size)
        return result_box

    def FreeResult(self):
        my_class_free_result = self.my_class_dll.FreeResult
        my_class_free_result.argtypes = [ctypes.POINTER(Detection), ctypes.POINTER(ctypes.c_size_t)]
        my_class_free_result.restype = None
        my_class_free_result(self.result_array, self.result_size)

    def Array2Box(self, result_array, result_size):
        result_box = []
        for i in range(result_size):
            box_dict = {}
            stru_box = result_array[i]
            clas_id = stru_box.clas_id
            conf = stru_box.conf
            box_array = stru_box.box_array
            box = box_array[0:4]
            box_dict["clas_id"] = clas_id
            box_dict["conf"] = conf
            box_dict["box"] = box

            result_box.append(box_dict)

        return result_box

    def DrawImg(self, img, result_box):
        self.draw_box = False  # 画轮廓
        # 循环多个类的轮廓
        for box_dict in result_box:
            clas_id = box_dict["clas_id"]
            label = self.class_name[clas_id]
            conf = round(box_dict["conf"], 3)
            box = box_dict["box"]

            x, y, w, h = box

            cv2.rectangle(img, (x, y), (x + w, y + h), color=self.color_list[clas_id], thickness=1)

            # 写类别
            # cv2.putText(img, label + f"_{area}", (x, y - 7), cv2.FONT_HERSHEY_SIMPLEX, 1, self.color_list[clas_id],3)
            cv2.putText(img, label + f"_{conf}", (x, y - 7), cv2.FONT_HERSHEY_SIMPLEX, 1, self.color_list[clas_id], 3)

        return img

    def SetColor(self, num_classes):
        arr_b = [0, 255, 128]
        arr_g = [255, 128, 0]
        arr_r = [128, 0, 255]
        size_channel = 3
        color_list = []
        for i in range(num_classes):
            c = i % size_channel
            ab = int(i / size_channel)
            b = ab % 10
            a = int(ab / 10)
            color = [arr_b[a], arr_g[b], arr_r[c]]
            color_list.append(color)
        return color_list


if __name__ == '__main__':
    imgs_path = r"../images"
    # dll_path = r"../lib/yolov5_onnx2dll_nm.dll"
    dll_path = r"D:\my_program\bushu\all_bushu\yolov5_detect\yolov5_onnxruntime_cpp\cmake-build-release\onnx2dll_nm/yolov5_onnx2dll_nm.dll"
    onnx_path = r"../weights/best_20230407.onnx"
    class_name = ["ng"]
    conf_threshold = 0.3

    my_class = MyClass(onnx_path, dll_path, class_name, conf_threshold)

    for img_name in os.listdir(imgs_path):
        print(img_name)
        img_path = os.path.join(imgs_path, img_name)
        img = cv2.imdecode(np.fromfile(img_path, np.uint8), 1)

        reslult = my_class.OutResult(img)
        print(len(reslult))
        my_class.DrawImg(img, reslult)

        # start_t = time.time()
        # for i in range(100):
        #     reslult = my_class.OutResult(img)
        #     my_class.DrawImg(img, reslult)
        #     my_class.FreeResult()
        # end_t = time.time()
        # print("detect_t:",(end_t-start_t)/100)

        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.show()

        # 释放堆区数组
        my_class.FreeResult()
        print("======================")
